import os
import shutil
from pcdet.utils import box_utils, calibration_kitti
import numpy as np

calib_file = './calib.txt'
calib = calibration_kitti.Calibration(calib_file)

load_dir = './label_big_range/'
save_dir = './label_2/'

for label in os.listdir(load_dir):
    ori_labelpath = os.path.join(load_dir, label)
    with open(ori_labelpath, encoding='utf-8') as f:
        label_buf = []
        objs = f.readlines()
        for obj in objs:
            tmp = obj.split(' ')
            cam_h, cam_w, cam_l = np.float32(tmp[8]), np.float32(tmp[9]), np.float32(tmp[10])
            cam_x, cam_y, cam_z = np.float32(tmp[11]), np.float32(tmp[12]), np.float32(tmp[13])
            cam_r = np.float32(tmp[14])
            boxes3d_camera = np.array([cam_x, cam_y, cam_z, cam_l, cam_h, cam_w, cam_r])
            boxes3d_lidar = box_utils.boxes3d_kitti_camera_to_lidar(boxes3d_camera.reshape(-1, 7), calib)
            x, y, z = boxes3d_lidar[0][0], boxes3d_lidar[0][1], boxes3d_lidar[0][2]
            if ((y < 40) and (y > -40)) and ((x > -40) and (x < 70)):
                label_buf.append(tmp)

    savefile = os.path.join(save_dir, label)
    file = open(savefile, 'w')
    for obj in label_buf:
        line = obj[0] + ' ' + obj[1] + ' ' + obj[2] + ' ' + obj[3] + ' ' + obj[4] + ' ' + obj[5] + ' ' + obj[6] + ' ' + \
               obj[7] + ' ' + obj[8] + ' ' + obj[9] + ' ' + obj[10] + ' ' + obj[11] + ' ' + obj[12] + ' ' + \
               obj[13] + ' ' + obj[14]
        file.write(str(line))
    file.close


